load("//xla:xla.default.bzl", "xla_cc_test", "xla_internal")
load("//xla/service/cpu:build_defs.bzl", "runtime_copts")
load("//xla/tsl:tsl.bzl", "if_google", "if_windows", "internal_visibility")
load("//xla/tsl:tsl.default.bzl", "filegroup")
load("//xla/tsl/platform:build_config.bzl", "tf_proto_library")
load("//xla/tsl/platform:rules_cc.bzl", "cc_library")

package(
    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
    default_visibility = internal_visibility([":friends"]),
    licenses = ["notice"],
)

package_group(
    name = "friends",
    includes = [
        "//xla:friends",
    ],
)

filegroup(
    name = "runtime_srcs",
    srcs = [
        "convolution_thunk_f16.cc",
        "convolution_thunk_f32.cc",
        "rng_state_lib.cc",
    ],
    visibility = internal_visibility([":friends"]),
)

filegroup(
    name = "runtime_hdrs",
    srcs = [
        "convolution_thunk_internal.h",
        "kernel_c_api.h",
        "rng_state_lib.h",
        "work_queue.h",
    ],
    visibility = internal_visibility([":friends"]),
)

cc_library(
    name = "buffer_allocations",
    hdrs = ["buffer_allocations.h"],
    deps = [
        "//xla:util",
        "//xla/service:buffer_assignment",
        "//xla/service:maybe_owning_device_memory",
        "//xla/stream_executor:device_memory",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/types:span",
    ],
)

xla_cc_test(
    name = "buffer_allocations_test",
    srcs = ["buffer_allocations_test.cc"],
    deps = [
        ":buffer_allocations",
        ":thunk_testlib",
        "//xla:literal_util",
        "//xla/service:buffer_assignment",
        "//xla/stream_executor:device_memory",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/platform:test",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "thread_pool_task_runner",
    hdrs = ["thread_pool_task_runner.h"],
    deps = [
        ":thunk",
        "@com_google_absl//absl/base:core_headers",
        "@eigen_archive//:eigen3",
    ],
)

cc_library(
    name = "kernel_c_api",
    hdrs = ["kernel_c_api.h"],
)

cc_library(
    name = "kernel",
    srcs = ["kernel.cc"],
    hdrs = ["kernel.h"],
    deps = [
        ":kernel_c_api",
        ":work_queue",
        "//xla:util",
        "//xla/runtime:work_group",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:launch_dim",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:logging",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/types:span",
        "@eigen_archive//:eigen3",
    ],
)

xla_cc_test(
    name = "kernel_test",
    srcs = ["kernel_test.cc"],
    deps = [
        ":kernel",
        ":kernel_c_api",
        "//xla/runtime:work_group",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/lib/core:status_test_util",
        "//xla/tsl/platform:env",
        "//xla/tsl/platform:test",
        "//xla/tsl/platform:test_benchmark",
        "//xla/tsl/platform:test_main",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/types:span",
        "@com_google_googletest//:gtest",
        "@eigen_archive//:eigen3",
        "@tsl//tsl/platform:platform_port",
    ],
)

cc_library(
    name = "function_library",
    srcs = ["function_library.cc"],
    hdrs = ["function_library.h"],
    deps = [
        ":kernel_c_api",
        "//xla/tsl/lib/gtl:int_type",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:string_view",
        "@tsl//tsl/platform:statusor",
    ],
)

tf_proto_library(
    name = "thunk_proto",
    srcs = ["thunk.proto"],
    create_grpc_library = True,
    make_default_target_header_only = True,
    protodeps = [
        "//xla/backends/cpu:xnnpack_config_proto",
        "//xla/service:buffer_assignment_proto",
        "//xla:xla_data_proto",
        "//xla/service:hlo_proto",
    ] + if_google(["@com_google_protobuf//:any"]),
    visibility = ["//visibility:public"],
)

cc_library(
    name = "thunk",
    srcs = ["thunk.cc"],
    hdrs = ["thunk.h"],
    deps = [
        ":buffer_allocations",
        ":function_library",
        ":xfeed_manager",
        "//xla:executable_run_options",
        "//xla/backends/cpu/collectives:cpu_collectives",
        "//xla/backends/cpu/collectives:in_process_collectives",
        "//xla/backends/cpu/runtime/xnnpack:xnn_interop",
        "//xla/backends/cpu/runtime/xnnpack:xnn_threadpool",
        "//xla/ffi:execution_context",
        "//xla/runtime:buffer_use",
        "//xla/runtime:resource_use",
        "//xla/service:global_device_id",
        "//xla/service/cpu:cpu_executable_run_options",
        "//xla/stream_executor:stream",
        "//xla/stream_executor:stream_executor_h",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:logging",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base:no_destructor",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/functional:function_ref",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:string_view",
        "@tsl//tsl/profiler/lib:traceme",
        "@tsl//tsl/profiler/lib:traceme_encode",
    ],
)

cc_library(
    name = "thunk_testlib",
    testonly = 1,
    srcs = ["thunk_testlib.cc"],
    hdrs = ["thunk_testlib.h"],
    deps = [
        ":buffer_allocations",
        ":thunk",
        "//xla:literal",
        "//xla/runtime:buffer_use",
        "//xla/runtime:resource_use",
        "//xla/service:buffer_assignment",
        "//xla/stream_executor:device_memory",
        "//xla/tsl/concurrency:async_value",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/types:span",
    ],
)

xla_cc_test(
    name = "thunk_test",
    srcs = ["thunk_test.cc"],
    deps = [
        ":thunk",
        "//xla:executable_run_options",
        "//xla/backends/cpu/collectives:cpu_collectives",
        "//xla/service/cpu:cpu_executable_run_options",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/platform:test",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "thunk_executor",
    srcs = ["thunk_executor.cc"],
    hdrs = ["thunk_executor.h"],
    local_defines = if_windows(["_ENABLE_EXTENDED_ALIGNED_STORAGE"]),
    deps = [
        ":thunk",
        "//xla/runtime:buffer_use",
        "//xla/runtime:execution_graph",
        "//xla/runtime:resource_use",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:env",
        "//xla/tsl/platform:logging",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:fixed_array",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
        "@tsl//tsl/platform:numbers",
        "@tsl//tsl/profiler/lib:connected_traceme",
        "@tsl//tsl/profiler/lib:context_types_hdrs",
        "@tsl//tsl/profiler/lib:traceme",
    ] + xla_internal(["service:execution_graph_visualizer_google"]),
)

xla_cc_test(
    name = "thunk_executor_test",
    srcs = ["thunk_executor_test.cc"],
    deps = [
        ":buffer_allocations",
        ":thread_pool_task_runner",
        ":thunk",
        ":thunk_executor",
        ":thunk_testlib",
        "//xla:literal",
        "//xla:literal_util",
        "//xla/runtime:buffer_use",
        "//xla/runtime:resource_use",
        "//xla/service:buffer_assignment",
        "//xla/stream_executor:device_memory",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:env",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:logging",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/platform:test",
        "//xla/tsl/platform:test_benchmark",
        "//xla/tsl/platform:test_main",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/types:span",
        "@eigen_archive//:eigen3",
    ],
)

cc_library(
    name = "call_thunk",
    srcs = ["call_thunk.cc"],
    hdrs = ["call_thunk.h"],
    deps = [
        ":thunk",
        ":thunk_executor",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/memory",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
    ],
)

cc_library(
    name = "conditional_thunk",
    srcs = ["conditional_thunk.cc"],
    hdrs = ["conditional_thunk.h"],
    deps = [
        ":thunk",
        ":thunk_executor",
        "//xla:util",
        "//xla/runtime:buffer_use",
        "//xla/service:buffer_assignment",
        "//xla/stream_executor:device_memory",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/memory",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:str_format",
    ],
)

xla_cc_test(
    name = "conditional_thunk_test",
    srcs = ["conditional_thunk_test.cc"],
    deps = [
        ":conditional_thunk",
        ":thunk",
        ":thunk_testlib",
        "//xla/runtime:buffer_use",
        "//xla/runtime:resource_use",
        "//xla/service:buffer_assignment",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/platform:test",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "all_gather_thunk",
    srcs = ["all_gather_thunk.cc"],
    hdrs = ["all_gather_thunk.h"],
    deps = [
        ":collective_thunk",
        ":thunk",
        "//xla:shape_util",
        "//xla/backends/cpu/collectives:cpu_collectives",
        "//xla/core/collectives:communicator",
        "//xla/service:collective_ops_utils",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/memory",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
    ],
)

cc_library(
    name = "convolution_lib",
    srcs = ["convolution_lib.cc"],
    hdrs = ["convolution_lib.h"],
    deps = [
        "//xla:shape_util",
        "//xla:status_macros",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/runtime:buffer_use",
        "//xla/service:buffer_assignment",
        "//xla/tsl/platform:errors",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/types:span",
    ],
)

cc_library(
    name = "convolution_thunk_internal",
    srcs = [
        "convolution_thunk_f16.cc",
        "convolution_thunk_f32.cc",
    ],
    hdrs = ["convolution_thunk_internal.h"],
    copts = runtime_copts(),
    visibility = internal_visibility([":friends"]),
    deps = [
        ":work_queue",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/framework/contraction:eigen_contraction_kernel",
        "//xla/tsl/framework/convolution:eigen_helpers",
        "//xla/tsl/platform:logging",
        "@eigen_archive//:eigen3",
    ],
)

cc_library(
    name = "convolution_thunk",
    srcs = ["convolution_thunk.cc"],
    hdrs = ["convolution_thunk.h"],
    copts = runtime_copts(),
    deps = [
        ":convolution_lib",
        ":convolution_thunk_internal",
        ":thunk",
        "//xla:executable_run_options",
        "//xla:shape_util",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/service:buffer_assignment",
        "//xla/stream_executor:device_memory",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/memory",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
        "@eigen_archive//:eigen3",
    ],
)

cc_library(
    name = "convolution_thunk_test_util",
    testonly = True,
    hdrs = ["convolution_thunk_test_util.h"],
    deps = [
        ":buffer_allocations",
        ":convolution_thunk",
        ":thunk_testlib",
        "//xla:literal",
        "//xla:literal_util",
        "//xla/service:buffer_assignment",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/types:span",
    ],
)

xla_cc_test(
    name = "convolution_thunk_test",
    srcs = ["convolution_thunk_test.cc"],
    deps = [
        ":buffer_allocations",
        ":convolution_thunk",
        ":convolution_thunk_test_util",
        ":thunk",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/platform:test",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/types:span",
        "@com_google_googletest//:gtest_main",
        "@eigen_archive//:eigen3",
    ],
)

cc_library(
    name = "all_reduce_thunk",
    srcs = ["all_reduce_thunk.cc"],
    hdrs = ["all_reduce_thunk.h"],
    deps = [
        ":collective_thunk",
        ":thunk",
        "//xla:shape_util",
        "//xla:util",
        "//xla/backends/cpu/collectives:cpu_collectives",
        "//xla/core/collectives:communicator",
        "//xla/service:buffer_assignment",
        "//xla/service:collective_ops_utils",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/memory",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
    ],
)

cc_library(
    name = "all_to_all_thunk",
    srcs = ["all_to_all_thunk.cc"],
    hdrs = ["all_to_all_thunk.h"],
    deps = [
        ":collective_thunk",
        ":thunk",
        "//xla:shape_util",
        "//xla:xla_data_proto_cc",
        "//xla/backends/cpu/collectives:cpu_collectives",
        "//xla/core/collectives:communicator",
        "//xla/service:buffer_assignment",
        "//xla/service:collective_ops_utils",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:logging",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/memory",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
    ],
)

cc_library(
    name = "reduce_scatter_thunk",
    srcs = ["reduce_scatter_thunk.cc"],
    hdrs = ["reduce_scatter_thunk.h"],
    deps = [
        ":collective_thunk",
        ":thunk",
        "//xla:shape_util",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/backends/cpu/collectives:cpu_collectives",
        "//xla/core/collectives:communicator",
        "//xla/service:collective_ops_utils",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/memory",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
    ],
)

cc_library(
    name = "collective_permute_thunk",
    srcs = ["collective_permute_thunk.cc"],
    hdrs = ["collective_permute_thunk.h"],
    deps = [
        ":collective_thunk",
        ":thunk",
        "//xla:shape_util",
        "//xla:status_macros",
        "//xla:xla_data_proto_cc",
        "//xla/backends/cpu/collectives:cpu_collectives",
        "//xla/core/collectives:communicator",
        "//xla/core/collectives:rank_id",
        "//xla/service:buffer_assignment",
        "//xla/service:collective_ops_utils",
        "//xla/service:computation_placer",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/memory",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/types:span",
    ],
)

cc_library(
    name = "collective_thunk",
    srcs = ["collective_thunk.cc"],
    hdrs = ["collective_thunk.h"],
    deps = [
        ":thunk",
        "//xla:shape_util",
        "//xla:status_macros",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/backends/cpu/collectives:cpu_clique_key",
        "//xla/backends/cpu/collectives:cpu_cliques",
        "//xla/backends/cpu/collectives:cpu_collectives",
        "//xla/core/collectives:communicator",
        "//xla/core/collectives:rank_id",
        "//xla/runtime:buffer_use",
        "//xla/runtime:resource_use",
        "//xla/service:buffer_assignment",
        "//xla/service:collective_ops_utils",
        "//xla/service:computation_placer",
        "//xla/service:global_device_id",
        "//xla/service:hlo_proto_cc",
        "//xla/stream_executor:device_memory",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:logging",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/functional:any_invocable",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/time",
        "@com_google_absl//absl/types:span",
    ],
)

cc_library(
    name = "copy_thunk",
    srcs = ["copy_thunk.cc"],
    hdrs = ["copy_thunk.h"],
    deps = [
        ":buffer_allocations",
        ":thunk",
        "//xla:shape_util",
        "//xla:util",
        "//xla/pjrt:transpose",
        "//xla/runtime:buffer_use",
        "//xla/service:buffer_assignment",
        "//xla/stream_executor:device_memory",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/memory",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
        "@eigen_archive//:eigen3",
    ],
)

xla_cc_test(
    name = "copy_thunk_test",
    srcs = ["copy_thunk_test.cc"],
    deps = [
        ":buffer_allocations",
        ":copy_thunk",
        ":thunk",
        ":thunk_testlib",
        "//xla:literal_util",
        "//xla:shape_util",
        "//xla:xla_data_proto_cc",
        "//xla/service:buffer_assignment",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/platform:test",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "custom_call_thunk",
    srcs = ["custom_call_thunk.cc"],
    hdrs = ["custom_call_thunk.h"],
    deps = [
        ":thunk",
        "//xla:shape_util",
        "//xla:util",
        "//xla/ffi:attribute_map",
        "//xla/ffi:call_frame",
        "//xla/ffi:execution_state",
        "//xla/ffi:ffi_api",
        "//xla/ffi/api:c_api",
        "//xla/runtime:buffer_use",
        "//xla/runtime:object_pool",
        "//xla/service:buffer_assignment",
        "//xla/service:custom_call_status",
        "//xla/service:custom_call_status_internal",
        "//xla/service:custom_call_target_registry",
        "//xla/service:hlo_proto_cc",
        "//xla/stream_executor:device_memory",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/base:dynamic_annotations",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/memory",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/strings:string_view",
        "@eigen_archive//:eigen3",
        "@llvm-project//mlir:AsmParser",
        "@llvm-project//mlir:IR",
        "@llvm-project//mlir:Support",
    ],
)

cc_library(
    name = "dot_lib",
    srcs = ["dot_lib.cc"],
    hdrs = ["dot_lib.h"],
    deps = [
        "//xla:shape_util",
        "//xla:status_macros",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/runtime:buffer_use",
        "//xla/service:buffer_assignment",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/types:span",
    ],
)

cc_library(
    name = "dot_thunk",
    srcs = [
        "dot_thunk.cc",
        "dot_thunk_c128.cc",
        "dot_thunk_c64.cc",
        "dot_thunk_f16.cc",
        "dot_thunk_f32.cc",
        "dot_thunk_f64.cc",
        "dot_thunk_s32.cc",
        "dot_thunk_s8.cc",
    ],
    hdrs = ["dot_thunk.h"],
    deps = [
        ":dot_lib",
        ":thunk",
        "//xla:shape_util",
        "//xla:types",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/service:buffer_assignment",
        "//xla/stream_executor:device_memory",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/framework/contraction:eigen_contraction_kernel",
        "//xla/tsl/platform:logging",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/functional:any_invocable",
        "@com_google_absl//absl/memory",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:str_format",
        "@eigen_archive//:eigen3",
    ],
)

xla_cc_test(
    name = "dot_thunk_test",
    srcs = ["dot_thunk_test.cc"],
    deps = [
        ":buffer_allocations",
        ":dot_thunk",
        ":thunk",
        ":thunk_testlib",
        "//xla:literal_util",
        "//xla:shape_util",
        "//xla:xla_data_proto_cc",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:env",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/platform:test",
        "@com_google_absl//absl/strings",
        "@eigen_archive//:eigen3",
        "@tsl//tsl/platform:test_main",
    ],
)

cc_library(
    name = "outfeed_thunk",
    srcs = ["outfeed_thunk.cc"],
    hdrs = ["outfeed_thunk.h"],
    deps = [
        ":thunk",
        ":xfeed_manager",
        "//xla:shape_util",
        "//xla:util",
        "//xla/runtime:buffer_use",
        "//xla/runtime:resource_use",
        "//xla/service:buffer_assignment",
        "//xla/stream_executor:device_memory",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/memory",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/types:span",
    ],
)

xla_cc_test(
    name = "outfeed_thunk_test",
    srcs = ["outfeed_thunk_test.cc"],
    deps = [
        ":outfeed_thunk",
        ":thunk",
        "//xla:shape_util",
        "//xla:xla_data_proto_cc",
        "//xla/runtime:buffer_use",
        "//xla/runtime:resource_use",
        "//xla/service:buffer_assignment",
        "@com_google_googletest//:gtest_main",
        "@tsl//tsl/platform:statusor",
        "@tsl//tsl/platform:test",
    ],
)

cc_library(
    name = "logical_id_thunk",
    srcs = ["logical_id_thunk.cc"],
    hdrs = ["logical_id_thunk.h"],
    deps = [
        ":thunk",
        "//xla:status_macros",
        "//xla/runtime:buffer_use",
        "//xla/service:buffer_assignment",
        "//xla/service:computation_placer_hdr",
        "//xla/service:global_device_id",
        "//xla/stream_executor:device_memory",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/memory",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
    ],
)

xla_cc_test(
    name = "logical_id_thunk_test",
    srcs = ["logical_id_thunk_test.cc"],
    deps = [
        ":buffer_allocations",
        ":logical_id_thunk",
        ":thunk",
        ":thunk_testlib",
        "//xla:executable_run_options",
        "//xla:literal_util",
        "//xla/service:buffer_assignment",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/platform:test",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "infeed_thunk",
    srcs = ["infeed_thunk.cc"],
    hdrs = ["infeed_thunk.h"],
    deps = [
        ":thunk",
        ":xfeed_manager",
        "//xla:shape_util",
        "//xla:util",
        "//xla/runtime:buffer_use",
        "//xla/runtime:resource_use",
        "//xla/service:buffer_assignment",
        "//xla/stream_executor:device_memory",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/memory",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/types:span",
    ],
)

xla_cc_test(
    name = "infeed_thunk_test",
    srcs = ["infeed_thunk_test.cc"],
    deps = [
        ":infeed_thunk",
        ":thunk",
        "//xla:shape_util",
        "//xla:xla_data_proto_cc",
        "//xla/runtime:buffer_use",
        "//xla/runtime:resource_use",
        "//xla/service:buffer_assignment",
        "@com_google_googletest//:gtest_main",
        "@tsl//tsl/platform:statusor",
        "@tsl//tsl/platform:test",
    ],
)

cc_library(
    name = "kernel_thunk",
    srcs = ["kernel_thunk.cc"],
    hdrs = ["kernel_thunk.h"],
    deps = [
        ":buffer_allocations",
        ":function_library",
        ":kernel",
        ":kernel_c_api",
        ":thunk",
        "//xla:util",
        "//xla/codegen:kernel_spec",
        "//xla/runtime:buffer_use",
        "//xla/runtime:work_group",
        "//xla/service:buffer_assignment",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:launch_dim",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/base",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_set",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/memory",
        "@com_google_absl//absl/numeric:bits",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_absl//absl/types:span",
        "@eigen_archive//:eigen3",
    ],
)

xla_cc_test(
    name = "kernel_thunk_test",
    srcs = ["kernel_thunk_test.cc"],
    deps = [
        ":buffer_allocations",
        ":function_library",
        ":kernel",
        ":kernel_c_api",
        ":kernel_thunk",
        ":thunk",
        ":thunk_testlib",
        "//xla:literal_util",
        "//xla/runtime:work_group",
        "//xla/service:buffer_assignment",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:launch_dim",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/platform:test",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "rng_state_lib",
    srcs = ["rng_state_lib.cc"],
    hdrs = ["rng_state_lib.h"],
    deps = [
        "@com_google_absl//absl/base:config",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/numeric:int128",
        "@com_google_absl//absl/synchronization",
    ],
)

xla_cc_test(
    name = "rng_state_lib_test",
    srcs = ["rng_state_lib_test.cc"],
    deps = [
        ":rng_state_lib",
        "@com_google_absl//absl/numeric:int128",
        "@com_google_googletest//:gtest_main",
    ],
)

cc_library(
    name = "rng_state_thunk",
    srcs = ["rng_state_thunk.cc"],
    hdrs = ["rng_state_thunk.h"],
    deps = [
        ":rng_state_lib",
        ":thunk",
        "//xla:util",
        "//xla/runtime:buffer_use",
        "//xla/service:buffer_assignment",
        "//xla/stream_executor:device_memory",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base:config",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/memory",
        "@com_google_absl//absl/numeric:int128",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/synchronization",
    ],
)

cc_library(
    name = "sort_thunk",
    srcs = ["sort_thunk.cc"],
    hdrs = ["sort_thunk.h"],
    deps = [
        ":function_library",
        ":thunk",
        "//xla:shape_util",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/runtime:buffer_use",
        "//xla/service:buffer_assignment",
        "//xla/stream_executor:device_memory",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:logging",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/base",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/base:dynamic_annotations",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/functional:any_invocable",
        "@com_google_absl//absl/memory",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/types:span",
    ],
)

xla_cc_test(
    name = "sort_thunk_test",
    srcs = ["sort_thunk_test.cc"],
    deps = [
        ":buffer_allocations",
        ":function_library",
        ":sort_thunk",
        ":thunk",
        ":thunk_testlib",
        "//xla:literal",
        "//xla:literal_util",
        "//xla:shape_util",
        "//xla:xla_data_proto_cc",
        "//xla/service:buffer_assignment",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:logging",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/platform:test",
        "//xla/tsl/platform:test_benchmark",
        "//xla/tsl/platform:test_main",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_absl//absl/types:span",
    ],
)

cc_library(
    name = "while_thunk",
    srcs = ["while_thunk.cc"],
    hdrs = ["while_thunk.h"],
    deps = [
        ":buffer_allocations",
        ":thunk",
        ":thunk_executor",
        "//xla/runtime:buffer_use",
        "//xla/service:buffer_assignment",
        "//xla/stream_executor:device_memory",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/memory",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:str_format",
    ],
)

xla_cc_test(
    name = "while_thunk_test",
    srcs = ["while_thunk_test.cc"],
    deps = [
        ":buffer_allocations",
        ":thunk",
        ":thunk_testlib",
        ":while_thunk",
        "//xla:literal_util",
        "//xla/runtime:buffer_use",
        "//xla/runtime:resource_use",
        "//xla/service:buffer_assignment",
        "//xla/stream_executor:device_memory",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:env",
        "//xla/tsl/platform:statusor",
        "//xla/tsl/platform:test",
        "@com_google_googletest//:gtest_main",
        "@eigen_archive//:eigen3",
    ],
)

cc_library(
    name = "fft_thunk",
    srcs = ["fft_thunk.cc"],
    hdrs = ["fft_thunk.h"],
    deps = [
        ":thunk",
        "//xla:shape_util",
        "//xla:status_macros",
        "//xla:xla_data_proto_cc",
        "//xla/runtime:buffer_use",
        "//xla/service:buffer_assignment",
        "//xla/stream_executor:device_memory",
        "//xla/stream_executor:stream_executor_h",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/container:inlined_vector",
        "@com_google_absl//absl/memory",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/types:span",
        "@ducc//:fft_wrapper",
        "@eigen_archive//:eigen3",
    ],
)

cc_library(
    name = "topk_thunk",
    srcs = ["topk_thunk.cc"],
    hdrs = ["topk_thunk.h"],
    deps = [
        ":thunk",
        "//xla/runtime:buffer_use",
        "//xla/service:buffer_assignment",
        "//xla/service/cpu:runtime_topk",
        "//xla/stream_executor:device_memory",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/memory",
        "@com_google_absl//absl/status:statusor",
    ],
)

cc_library(
    name = "serdes_base",
    hdrs = ["serdes_base.h"],
    deps = ["@com_google_absl//absl/status:statusor"],
)

cc_library(
    name = "thunk_proto_serdes",
    srcs = ["thunk_proto_serdes.cc"],
    hdrs = ["thunk_proto_serdes.h"],
    deps = [
        ":all_gather_thunk",
        ":all_reduce_thunk",
        ":all_to_all_thunk",
        ":call_thunk",
        ":collective_permute_thunk",
        ":collective_thunk",
        ":conditional_thunk",
        ":convolution_lib",
        ":convolution_thunk",
        ":copy_thunk",
        ":custom_call_thunk",
        ":dot_thunk",
        ":fft_thunk",
        ":infeed_thunk",
        ":kernel_thunk",
        ":logical_id_thunk",
        ":outfeed_thunk",
        ":reduce_scatter_thunk",
        ":rng_state_thunk",
        ":serdes_base",
        ":sort_thunk",
        ":thunk",
        ":thunk_proto_cc_impl",
        ":topk_thunk",
        ":while_thunk",
        "//xla:shape_util",
        "//xla:util",
        "//xla/backends/cpu:xnnpack_config_proto_cc",
        "//xla/backends/cpu/runtime/xnnpack:xnn_convolution_thunk",
        "//xla/backends/cpu/runtime/xnnpack:xnn_dot_thunk",
        "//xla/backends/cpu/runtime/xnnpack:xnn_fusion_thunk",
        "//xla/runtime:resource_use",
        "//xla/runtime:work_group",
        "//xla/service:buffer_assignment",
        "//xla/service:collective_ops_utils",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/container:flat_hash_set",
        "@com_google_absl//absl/log",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/strings:string_view",
        "@tsl//tsl/platform:casts",
    ],
)

cc_library(
    name = "work_queue",
    hdrs = ["work_queue.h"],
    deps = [
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:logging",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:fixed_array",
        "@com_google_absl//absl/status",
        "@eigen_archive//:eigen3",
    ],
)

xla_cc_test(
    name = "work_queue_test",
    srcs = ["work_queue_test.cc"],
    deps = [
        ":work_queue",
        "//xla/tsl/concurrency:async_value",
        "//xla/tsl/platform:env",
        "//xla/tsl/platform:test",
        "//xla/tsl/platform:test_benchmark",
        "//xla/tsl/platform:test_main",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/synchronization",
        "@eigen_archive//:eigen3",
    ],
)

xla_cc_test(
    name = "thunk_sequence_serdes_test",
    srcs = ["thunk_sequence_serdes_test.cc"],
    deps = [
        ":all_gather_thunk",
        ":all_reduce_thunk",
        ":all_to_all_thunk",
        ":call_thunk",
        ":collective_permute_thunk",
        ":collective_thunk",
        ":conditional_thunk",
        ":convolution_thunk",
        ":convolution_thunk_test_util",
        ":copy_thunk",
        ":custom_call_thunk",
        ":dot_thunk",
        ":fft_thunk",
        ":infeed_thunk",
        ":kernel",
        ":kernel_thunk",
        ":logical_id_thunk",
        ":outfeed_thunk",
        ":reduce_scatter_thunk",
        ":rng_state_thunk",
        ":serdes_base",
        ":sort_thunk",
        ":thunk",
        ":thunk_executor",
        ":thunk_proto_cc",
        ":thunk_proto_serdes",
        ":thunk_testlib",
        ":topk_thunk",
        ":while_thunk",
        "//xla:literal",
        "//xla:literal_util",
        "//xla:shape_util",
        "//xla:util",
        "//xla:xla_data_proto_cc",
        "//xla/backends/cpu/runtime/xnnpack:xnn_convolution_thunk",
        "//xla/backends/cpu/runtime/xnnpack:xnn_dot_thunk",
        "//xla/backends/cpu/runtime/xnnpack:xnn_fusion_thunk",
        "//xla/ffi",
        "//xla/ffi:ffi_api",
        "//xla/runtime:resource_use",
        "//xla/runtime:work_group",
        "//xla/service:buffer_assignment",
        "//xla/service:collective_ops_utils",
        "//xla/service:hlo_proto_cc",
        "//xla/stream_executor:launch_dim",
        "//xla/tsl/platform:errors",
        "//xla/tsl/platform:status",
        "//xla/tsl/platform:statusor",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/container:flat_hash_set",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/strings:string_view",
        "@com_google_absl//absl/types:span",
        "@com_google_googletest//:gtest_main",
        "@tsl//tsl/platform:casts",
    ],
)

cc_library(
    name = "xfeed_manager",
    srcs = ["xfeed_manager.cc"],
    hdrs = ["xfeed_manager.h"],
    deps = [
        "//xla:shape_util",
        "//xla:xla_data_proto_cc",
        "//xla/tsl/platform:logging",
        "@com_google_absl//absl/base:core_headers",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/status:statusor",
        "@com_google_absl//absl/synchronization",
        "@com_google_absl//absl/types:span",
    ],
)

xla_cc_test(
    name = "xfeed_manager_test",
    srcs = ["xfeed_manager_test.cc"],
    deps = [
        ":xfeed_manager",
        "//xla:shape_util",
        "//xla:xla_data_proto_cc",
        "//xla/tests:xla_internal_test_main",
        "//xla/tsl/lib/core:status_test_util",
        "//xla/tsl/platform:env",
        "//xla/tsl/platform:logging",
        "//xla/tsl/platform:test",
        "@com_google_absl//absl/log:check",
        "@com_google_absl//absl/status:statusor",
        "@com_google_googletest//:gtest",
    ],
)
